import pickle
import io
import dep_graph.keys
from bsddb3 import db

class PickledDatabase(object):
  """Handles pickling/unpickling objects into a bsddb3.db.DB.

  When used as a context manager, closes itself after the end of the context.
  
  Attributes:
    __db: Our bsddb3.db.DB.
  """
  class __Pickler(pickle.Pickler):
    def __init__(self, byteio, graph):
      super().__init__(byteio)
      self.graph = graph

    def persistent_id(self, obj):
      if obj == self.graph:
        return "Graph"
      else:
        return None

  class __Unpickler(pickle.Unpickler):
    def __init__(self, byteio, graph):
      super().__init__(byteio)
      self.graph = graph

    def persistent_load(self, val):
      if val == "Graph":
        return self.graph
      else:
        raise pickle.UnpicklingError("unsupported persistent object %s" % val)

  """Args:
    fname: The filename for the database.
    truncate: If True, empty the database out before doing anything.
  """
  def __init__(self, fname, truncate=False):
    self.__db = db.DB()
    flags = db.DB_TRUNCATE | db.DB_CREATE if truncate else db.DB_CREATE
    self.__db.open(fname, dbtype=db.DB_HASH, flags=flags)

  def write(self, key, value):
    """Writes a mapping into the database.

    Args:
      key: The key to set.
      value: The object to set that key to.
    """
    data = io.BytesIO()
    PickledDatabase.__Pickler(data, self).dump(value)
    self.__db[repr(key).encode('utf-8')] = data.getvalue()

  def lookup(self, key):
    """Gets a mapping from the database.

    Args:
      key: The key to get.

    Returns the value key maps to or None."""
    bytes_key = repr(key).encode('utf-8')
    if bytes_key in self.__db:
      return PickledDatabase.__Unpickler(io.BytesIO(self.__db[bytes_key]),
                                         self).load()

  def keys(self):
    """Yields all of the keys in the database."""
    return (dep_graph.keys.GenericGraphKey(key.decode('utf-8'))
        for key in self.__db.keys())

  def close(self):
    self.__db.close()

  def remove(self, key):
    """Removes a mapping from the database.

    Args:
      key: The key to remove.
    """
    try:
      del self.__db[repr(key).encode('utf-8')]
    except Exception as e:
      print("could not delete" + repr(key))
      raise e

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):  
    self.close()
